#!/usr/bin/python3
# coding=utf-8

# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# This file is a part of the CANN Open Software.
# Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# ======================================================================================================================
import os
import sys
import csv
import time
import logging
import subprocess
import shlex

import numpy as np

sys.path.append("../..")
from gen_quant_data import QuantMatmulGenData
from common_scripts.compare_data import compare_data
from common_scripts.exec_utils import get_file_work_dir, get_case_list, clear_file_cache, get_perf_task_duration

IS_BIAS = True
IS_TRANS_A = False
IS_TRANS_B = True
# int8_float16_dequant(int8 in float16 out for quant matmul)
DATA_TYPE_STR = "int8_float16_dequant"

logging.basicConfig(level=logging.INFO)


class ProcessParams:
    def __init__(self, case_name, m, n, k, quant_mode, is_perf, run_mode):
        self.case_name = case_name
        self.m = m
        self.n = n
        self.k = k
        self.quant_mode = quant_mode
        self.is_perf = is_perf
        self.run_mode = run_mode


def process_case(file_work_dir, process_params):
    case_name, m, n, k, quant_mode, is_perf, run_mode = process_params.case_name, process_params.m, process_params.n,\
        process_params.k, process_params.quant_mode, process_params.is_perf, process_params.run_mode
    logging.info("[INFO] start process case [%s]" % (case_name))
    logging.info("[INFO] IS_PERF is set [%d]" % (is_perf))
    clear_file_cache(file_work_dir)

    matmul_gen_data = QuantMatmulGenData(m, n, k, quant_mode, IS_TRANS_A, IS_TRANS_B, IS_BIAS, DATA_TYPE_STR)
    params_str = f"{m} {n} {k} {quant_mode}"
    matmul_gen_data.gen_golden_data(file_work_dir)
    if is_perf:
        cmd = f"msprof op --application=\"./ascendc_matmul_quant_bbit {params_str}\" --output=./prof_out"
    elif run_mode == "sim":
        cmd = f"msprof op simulator --application=\"./ascendc_matmul_quant_bbit {params_str}\" --output=./sim_out"
    else:
        cmd = f"./ascendc_matmul_quant_bbit {params_str}"
    subprocess.run(shlex.split(cmd))
    if is_perf:
        wrong_num = -1
    else:
        logging.info("[INFO] compare data case[%s]" % (case_name))
        wrong_num = compare_data(file_work_dir, DATA_TYPE_STR)
    res_data = []
    res_data.append(case_name)
    res_data.append(wrong_num)
    res_data.append(m * n)
    if wrong_num == -1:
        res_data.append("None")
    elif wrong_num / (m * n) > 0.001:
        res_data.append("Fail")
    else:
        res_data.append("Success")
    if is_perf:
        task_duration = get_perf_task_duration("./prof_out")
        res_data.append(task_duration)
    return res_data


def main():
    args_len = len(sys.argv) - 1
    if args_len != 2:
        logging.info("[ERROR] exec_test input params error!")
        return -1

    file_work_dir = get_file_work_dir()
    if not os.path.exists(file_work_dir):
        logging.info("[ERROR] file path %s not exist!" % (file_work_dir))
        return -1

    is_perf = False
    title_list = ["case_name", "wrong_num", "total_num", "result"]
    if sys.argv[2] == "perf":
        is_perf = True
        title_list.append("task_duration")

    case_list = get_case_list()
    res_list = [title_list]
    run_mode = sys.argv[1]
    for is_process, case_name, m, n, k, quant_mode in case_list:
        if is_process == 1:
            process_params = ProcessParams(case_name, m, n, k, quant_mode, is_perf, run_mode)
            res_data = process_case(file_work_dir, process_params)
            res_list.append(res_data)

    timestamp = time.time()
    result_file_name = "result_" + str(timestamp) + ".csv"
    with open(os.path.join(file_work_dir, "output", result_file_name), 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerows(res_list)

    logging.info("---------------RESULT---------------")
    for res in res_list:
        logging.info(res)
    return 0


if __name__ == "__main__":
    main()
